#!/bin/bash
set -e

echo "===== 增强版 pgBackRest 服务器B 恢复环境初始化 ====="

# 加载环境变量
if [ -f "/.env" ]; then
  echo "从根目录加载环境变量..."
  source /.env
elif [ -f ".env" ]; then
  echo "从当前目录加载环境变量..."
  source .env
elif [ -f "../.env" ]; then
  echo "从上级目录加载环境变量..."
  source ../.env
fi

# 显示加载的环境变量
echo "环境变量:"
echo "  PG_HOST_A: ${PG_HOST_A:-未设置}"
echo "  PG_HOST_B: ${PG_HOST_B:-未设置}"
echo "  SERVER_TYPE: ${SERVER_TYPE:-B}"  # 服务器B默认值

# 设置备份目录 - 优先使用.env中的值
BACKUP_DIR=${BACKUP_DIR:-/var/lib/pgbackrest}
LOG_LEVEL=${LOG_LEVEL:-info}
STANZA_NAME=${STANZA_NAME:-poddb}   # 可配置的stanza名称
OPERATION=${OPERATION:-init}        # 可选操作：init(初始化)或restore(恢复)

echo "配置信息:"
echo "  备份目录: ${BACKUP_DIR}"
echo "  日志级别: ${LOG_LEVEL}"
echo "  Stanza名称: ${STANZA_NAME}"
echo "  操作模式: ${OPERATION}"

# 添加超时函数
run_with_timeout() {
  local timeout=$1
  local command=$2
  local message=$3

  echo "执行命令(${timeout}秒超时): $command"

  # 使用timeout命令运行，并捕获输出
  timeout ${timeout}s bash -c "$command" > /tmp/cmd_output.log 2>&1
  local exit_code=$?

  # 显示输出
  cat /tmp/cmd_output.log

  # 检查超时
  if [ $exit_code -eq 124 ]; then
    echo "警告: 命令执行超时 (${timeout}秒)"
    echo "$message"
    return 1
  elif [ $exit_code -ne 0 ]; then
    echo "警告: 命令执行失败，退出码 $exit_code"
    return $exit_code
  fi

  return 0
}

# 检测服务器类型 - 优先使用环境变量SERVER_TYPE
SERVER_TYPE=${SERVER_TYPE:-"B"}  # 强制设置为服务器B
echo "服务器类型: 服务器${SERVER_TYPE}"

# 设置相关变量
CONTAINER_NAME="pgbackrest-${SERVER_TYPE}"
PG_CONTAINER_NAME="postgres-${SERVER_TYPE}"

# 设置PG_HOST - 优先使用环境变量
PG_HOST=${PG_HOST:-${PG_HOST_B:-$PG_CONTAINER_NAME}}

echo "使用容器名称: ${CONTAINER_NAME}"
echo "连接数据库主机: ${PG_HOST}"

# 设置PostgreSQL连接信息 - 优先使用.env中的值
PG_USER=${PG_USER:-postgres}
PG_PASSWORD=${PG_PASSWORD:-thingbelt2024}
PG_DB=${PG_DATABASE:-postgres}
PG_PORT=${PG_PORT:-5432}

echo "PostgreSQL连接信息:"
echo "  主机: ${PG_HOST}"
echo "  端口: ${PG_PORT}"
echo "  用户: ${PG_USER}"
echo "  数据库: ${PG_DB}"

# 创建必要的目录
echo "1. 创建必要的目录..."
mkdir -p /var/log/pgbackrest
chmod 777 /var/log/pgbackrest

# 创建并准备备份目录结构
mkdir -p ${BACKUP_DIR}
chmod 777 ${BACKUP_DIR}
mkdir -p ${BACKUP_DIR}/backup
chmod 777 ${BACKUP_DIR}/backup
mkdir -p ${BACKUP_DIR}/archive
chmod 777 ${BACKUP_DIR}/archive

# 为归档准备目录结构
mkdir -p ${BACKUP_DIR}/archive/${STANZA_NAME}
chmod 777 ${BACKUP_DIR}/archive/${STANZA_NAME}

echo "  备份目录结构已准备好: ${BACKUP_DIR}"
ls -la ${BACKUP_DIR}

# 确保容器运行 - 不需要检查容器状态，因为我们已经在容器内部
echo "2. 确认服务状态..."
echo "  当前运行在: ${CONTAINER_NAME}"

# 初始化日志目录
echo "3. 初始化日志目录..."
mkdir -p /var/log/pgbackrest && chmod 777 /var/log/pgbackrest

# 清空现有锁文件
echo "4. 清除可能存在的锁文件..."
mkdir -p /tmp/pgbackrest
rm -f /tmp/pgbackrest/*.lock

# 测试PostgreSQL连接
echo "5. 测试连接到PostgreSQL..."
PG_CONNECTED=0
if command -v pg_isready >/dev/null 2>&1; then
  pg_isready -h ${PG_HOST} -p ${PG_PORT} -U ${PG_USER} && PG_CONNECTED=1 || {
    echo "警告：无法连接到PostgreSQL，请确保数据库正在运行";
    echo "尝试等待PostgreSQL启动 (30秒)...";
    sleep 30;
    pg_isready -h ${PG_HOST} -p ${PG_PORT} -U ${PG_USER} && PG_CONNECTED=1 || {
      echo "警告：PostgreSQL仍然无法连接，但会继续尝试";
    }
  }
else
  echo "警告：找不到pg_isready命令，跳过PostgreSQL连接测试"
  echo "请确保安装了PostgreSQL客户端工具"
fi

# 自动检测PostgreSQL数据目录
PG_DATA_DIR=${PG_DATA_DIR:-"/var/lib/postgresql/data"}  # 允许通过环境变量覆盖
if [ $PG_CONNECTED -eq 1 ]; then
  echo "6. 尝试检测PostgreSQL数据目录位置..."
  export PGPASSWORD="${PG_PASSWORD}"

  # 尝试通过查询获取data_directory
  DB_DATA_DIR=$(psql -h ${PG_HOST} -p ${PG_PORT} -U ${PG_USER} -d ${PG_DB} -t -c "SHOW data_directory;" 2>/dev/null | xargs)
  if [ -n "${DB_DATA_DIR}" ]; then
    echo "  从PostgreSQL检测到数据目录: ${DB_DATA_DIR}"
    PG_DATA_DIR="${DB_DATA_DIR}"
  else
    echo "  无法从PostgreSQL查询获取数据目录，使用默认值或环境变量定义的值..."
  fi

  echo "  选择的PostgreSQL数据目录: ${PG_DATA_DIR}"

  # 尝试检测实际控制文件位置
  echo "  尝试验证pg_control文件位置..."

  # 首先尝试检查本地路径
  if test -f "${PG_DATA_DIR}/global/pg_control" 2>/dev/null; then
    echo "  √ 已在本地确认pg_control文件存在: ${PG_DATA_DIR}/global/pg_control"
    # 如果可以在本地找到，设置本地模式
    PG_DATA_MODE="local"
  else
    echo "  ! 在本地未找到pg_control文件。尝试通过docker inspect查找数据卷..."

    # 尝试通过Docker inspect查找
    if command -v docker > /dev/null 2>&1; then
      PG_CONTAINER_PATH=$(docker inspect ${PG_CONTAINER_NAME} 2>/dev/null | grep -A 10 "Volumes" | grep "/var/lib/postgresql/data" | head -1 | awk -F'"' '{print $2}')
      if [ -n "${PG_CONTAINER_PATH}" ]; then
        echo "  从Docker获取到数据目录路径: ${PG_CONTAINER_PATH}"
        PG_DATA_DIR="${PG_CONTAINER_PATH}"
        if test -f "${PG_DATA_DIR}/global/pg_control" 2>/dev/null; then
          echo "  √ 已在Docker卷中确认pg_control文件存在: ${PG_DATA_DIR}/global/pg_control"
          PG_DATA_MODE="docker"
        fi
      else
        echo "  无法从Docker获取数据目录路径"
      fi
    fi

    if [ -z "${PG_DATA_MODE}" ]; then
      echo "  ! 警告: 无法确认pg_control文件位置"
      echo "  假设PostgreSQL数据目录位于容器内，设置默认模式"
      PG_DATA_MODE="docker"
      echo "  将使用标准路径: ${PG_DATA_DIR}"
    fi
  fi
else
  echo "6. 无法连接到PostgreSQL，将使用默认或环境变量设置的数据目录: ${PG_DATA_DIR}"
  echo "  假设PostgreSQL数据目录位于容器内，设置默认模式"
  PG_DATA_MODE="docker"
fi

# 检查pgBackRest命令是否存在
if ! command -v pgbackrest >/dev/null 2>&1; then
  echo "错误：找不到pgbackrest命令，请确保pgBackRest已正确安装"
  echo "请在容器中安装pgBackRest：apt-get update && apt-get install -y pgbackrest"
  exit 1
fi

# 显示pgBackRest版本
echo "7. pgBackRest版本信息:"
pgbackrest version

# 配置pgBackRest
echo "8. 配置pgBackRest..."
mkdir -p /etc/pgbackrest
chmod 750 /etc/pgbackrest

# 创建pgbackrest.conf配置文件
echo "  创建配置文件..."
cat > /etc/pgbackrest/pgbackrest.conf << EOF
[global]
log-path=/var/log/pgbackrest
repo1-path=${BACKUP_DIR}
process-max=4
start-fast=y
delta=y
compress-level=6

# 以下选项禁用不必要的检查，提高成功率
archive-check=n
archive-mode=preserve
archive-timeout=300

[${STANZA_NAME}]
pg1-path=${PG_DATA_DIR}
EOF

# 根据数据目录访问模式，添加额外配置
if [ "${PG_DATA_MODE}" = "docker" ]; then
    # Docker模式的选项
    cat >> /etc/pgbackrest/pgbackrest.conf << EOF
archive-async=y
archive-push-queue-max=1GiB
archive-timeout=60
backup-standby=n
buffer-size=16MB
EOF
else
    # 本地模式的选项
    cat >> /etc/pgbackrest/pgbackrest.conf << EOF
archive-async=y
archive-push-queue-max=1GiB
archive-timeout=60
backup-standby=n
buffer-size=16MB
EOF
fi

echo "  配置文件内容:"
cat /etc/pgbackrest/pgbackrest.conf

# 确保PostgreSQL数据目录存在
echo "9. 确保数据目录存在..."
if [ -d "${PG_DATA_DIR}" ]; then
  echo "  √ 数据目录已存在: ${PG_DATA_DIR}"
else
  mkdir -p ${PG_DATA_DIR} || {
    echo "  警告：无法创建PostgreSQL数据目录，可能是因为目录在另一个容器中。";
    echo "  脚本将继续执行...";
  }
fi

# 设置环境变量为PostgreSQL设置
export PGHOST="${PG_HOST}"
export PGPORT="${PG_PORT}"
export PGUSER="${PG_USER}"
export PGPASSWORD="${PG_PASSWORD}"
export PGDATABASE="${PG_DB}"

# 根据操作模式选择执行流程
if [ "${OPERATION}" = "restore" ]; then
  # 恢复操作流程
  echo "10. 准备执行恢复操作..."

  # 检查PostgreSQL是否在运行
  echo "  检查PostgreSQL服务状态..."
  PG_RUNNING=false

  # 检查postmaster.pid文件
  if [ -f "${PG_DATA_DIR}/postmaster.pid" ]; then
    echo "  ⚠️ 检测到PostgreSQL正在运行！恢复操作需要PostgreSQL停止。"
    echo "  发现进程文件: ${PG_DATA_DIR}/postmaster.pid"

    # 尝试检查进程是否真的存在
    PID=$(head -1 "${PG_DATA_DIR}/postmaster.pid" 2>/dev/null || echo "")
    if [ -n "$PID" ] && ps -p $PID >/dev/null 2>&1; then
      echo "  确认PostgreSQL进程正在运行 (PID: $PID)"
      PG_RUNNING=true
    else
      echo "  进程ID存在但可能不在运行，尝试移除pid文件..."
      if [ -w "${PG_DATA_DIR}/postmaster.pid" ]; then
        rm -f "${PG_DATA_DIR}/postmaster.pid"
        echo "  已删除过时的pid文件"
      else
        echo "  ⚠️ 无权限删除pid文件，请手动停止PostgreSQL或删除pid文件"
        PG_RUNNING=true
      fi
    fi
  else
    echo "  ✅ 未检测到PostgreSQL运行，可以进行恢复操作"
  fi

  if [ "$PG_RUNNING" = true ]; then
    echo "  ❌ 错误: 恢复操作需要PostgreSQL停止。请执行以下操作:"
    echo "    1. 停止PostgreSQL服务:"
    echo "       docker stop postgres-${SERVER_TYPE}"
    echo "    2. 确认服务已停止后，重新运行此脚本:"
    echo "       docker exec -e OPERATION=restore pgbackrest-${SERVER_TYPE} /scripts/init-pgbackrest-server-b-enhanced.sh"
    echo ""
    echo "  或者您也可以手动设置FORCE_RESTORE=true来忽略此检查(不推荐):"
    echo "  docker exec -e OPERATION=restore -e FORCE_RESTORE=true pgbackrest-${SERVER_TYPE} /scripts/init-pgbackrest-server-b-enhanced.sh"

    # 如果设置了强制恢复，允许继续
    if [ "$FORCE_RESTORE" = "true" ]; then
      echo "  🔄 已设置强制恢复模式，将尝试继续..."
      echo "  ⚠️ 注意：这可能导致数据损坏，仅用于特殊情况！"
      if [ -f "${PG_DATA_DIR}/postmaster.pid" ] && [ -w "${PG_DATA_DIR}/postmaster.pid" ]; then
        echo "  尝试删除pid文件..."
        rm -f "${PG_DATA_DIR}/postmaster.pid"
      fi
    else
      exit 1
    fi
  fi

  # 检查备份目录是否有有效备份
  echo "  检查备份文件..."
  if [ -d "${BACKUP_DIR}/backup" ]; then
    BACKUP_COUNT=$(find "${BACKUP_DIR}/backup" -maxdepth 1 -type d | wc -l)
    if [ "$BACKUP_COUNT" -gt 1 ]; then
      echo "  ✅ 发现备份目录，共 $((BACKUP_COUNT-1)) 个备份"

      # 检查stanza是否已存在
      echo "  检查stanza状态..."
      if pgbackrest info --stanza=${STANZA_NAME} &>/dev/null; then
        echo "  ✅ Stanza '${STANZA_NAME}' 已存在，准备恢复"
      else
        echo "  ! Stanza '${STANZA_NAME}' 不存在，尝试创建..."
        pgbackrest --config=/etc/pgbackrest/pgbackrest.conf --stanza=${STANZA_NAME} --no-online stanza-create || {
          echo "  ! 创建stanza失败，尝试使用--force选项..."
          pgbackrest --config=/etc/pgbackrest/pgbackrest.conf --stanza=${STANZA_NAME} --no-online --force stanza-create || {
            echo "  ❌ 无法创建stanza，请检查配置和备份文件"
            exit 1
          }
        }
      fi

      # 备份原数据目录
      if [ -d "${PG_DATA_DIR}" ] && [ "$(ls -A ${PG_DATA_DIR} 2>/dev/null)" ]; then
        echo "  数据目录不为空，创建备份..."
        BACKUP_TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
        BACKUP_DATA_DIR="${PG_DATA_DIR}_backup_${BACKUP_TIMESTAMP}"
        mv "${PG_DATA_DIR}" "${BACKUP_DATA_DIR}" 2>/dev/null || {
          echo "  无法移动现有数据目录，尝试使用临时目录..."
          mkdir -p "/tmp/pg_data_backup_${BACKUP_TIMESTAMP}"
          cp -r "${PG_DATA_DIR}"/* "/tmp/pg_data_backup_${BACKUP_TIMESTAMP}/" 2>/dev/null || {
            echo "  ⚠️ 无法备份现有数据，请确保有足够权限"
          }
        }
        mkdir -p "${PG_DATA_DIR}"
        chmod 700 "${PG_DATA_DIR}"
      fi

      # 执行恢复操作
      echo "  准备执行数据恢复..."
      echo "  数据目录: ${PG_DATA_DIR}"
      echo "  操作将在5秒后开始，按Ctrl+C取消..."
      sleep 5

      # 构建恢复命令
      RESTORE_CMD="pgbackrest --config=/etc/pgbackrest/pgbackrest.conf --stanza=${STANZA_NAME} --log-level-console=${LOG_LEVEL} --pg1-path=${PG_DATA_DIR} --repo1-path=${BACKUP_DIR} restore"
      echo "  执行恢复命令: ${RESTORE_CMD}"

      ${RESTORE_CMD} || {
        echo "  ! 恢复失败，尝试添加--force选项..."
        ${RESTORE_CMD} --force || {
          echo "  ❌ 恢复操作失败，尝试更多debug选项..."
          ${RESTORE_CMD} --force --log-level-console=debug || {
            echo "  ❌ 恢复操作完全失败，请检查详细错误日志:"
            find /var/log/pgbackrest -type f -name "*.log" | xargs tail -n 50
            exit 1
          }
        }
      }

      echo "  ✅ 恢复操作完成"
      echo "  设置数据目录权限..."
      chmod 700 ${PG_DATA_DIR}
      chown -R postgres:postgres ${PG_DATA_DIR} 2>/dev/null || echo "  ⚠️ 无法修改权限，可能需要手动调整"

      # 添加恢复后优化部分
      echo "  执行恢复后优化..."

      # 1. 创建恢复配置文件告知PostgreSQL恢复已完成
      if [ -f "${PG_DATA_DIR}/recovery.signal" ]; then
        echo "  删除recovery.signal文件..."
        rm -f "${PG_DATA_DIR}/recovery.signal"
      fi

      # 2. 调整PostgreSQL配置，避免启动问题
      if [ -f "${PG_DATA_DIR}/postgresql.conf" ]; then
        echo "  优化PostgreSQL配置..."

        # 备份原配置
        cp ${PG_DATA_DIR}/postgresql.conf ${PG_DATA_DIR}/postgresql.conf.bak

        # 移除可能导致冲突的设置
        sed -i '/hot_standby =/d' ${PG_DATA_DIR}/postgresql.conf
        sed -i '/wal_level =/d' ${PG_DATA_DIR}/postgresql.conf

        # 添加安全启动参数
        echo "# 添加恢复后的优化参数" >> ${PG_DATA_DIR}/postgresql.conf
        echo "hot_standby = on" >> ${PG_DATA_DIR}/postgresql.conf
        echo "wal_level = replica" >> ${PG_DATA_DIR}/postgresql.conf
        echo "listen_addresses = '*'" >> ${PG_DATA_DIR}/postgresql.conf
      fi

      # 3. 确保pg_hba.conf允许访问
      if [ -f "${PG_DATA_DIR}/pg_hba.conf" ]; then
        echo "  确保pg_hba.conf配置正确..."

        # 备份原配置
        cp ${PG_DATA_DIR}/pg_hba.conf ${PG_DATA_DIR}/pg_hba.conf.bak

        # 检查pg_hba.conf是否包含必要条目，如果没有则添加
        if ! grep -q "host all all 0.0.0.0/0 md5" ${PG_DATA_DIR}/pg_hba.conf; then
          echo "# 允许所有IP通过密码访问" >> ${PG_DATA_DIR}/pg_hba.conf
          echo "host all all 0.0.0.0/0 md5" >> ${PG_DATA_DIR}/pg_hba.conf
        fi
      fi

      # 4. 删除可能导致启动问题的锁文件
      echo "  清理临时文件和锁..."
      rm -f ${PG_DATA_DIR}/postmaster.pid 2>/dev/null || true
      rm -f ${PG_DATA_DIR}/postgresql.trigger.* 2>/dev/null || true
      find ${PG_DATA_DIR} -name "*.tmp" -delete 2>/dev/null || true

      # 5. 调整关键文件权限
      echo "  设置关键文件权限..."
      chmod 600 ${PG_DATA_DIR}/postgresql.conf ${PG_DATA_DIR}/pg_hba.conf 2>/dev/null || true

      # 6. 如果存在.pgpass文件，确保权限正确
      if [ -f "${PG_DATA_DIR}/.pgpass" ]; then
        chmod 600 ${PG_DATA_DIR}/.pgpass
      fi

      # 记录恢复时间到日志文件，方便后续排查
      echo "恢复完成时间: $(date)" > ${PG_DATA_DIR}/restore_completed.log

      echo "  恢复后操作指南:"
      echo "  1. 确保PostgreSQL数据目录权限正确: chown -R postgres:postgres ${PG_DATA_DIR}"
      echo "  2. 启动PostgreSQL服务: docker start postgres-${SERVER_TYPE}"
      echo "  3. 检查PostgreSQL日志确认启动成功: docker logs postgres-${SERVER_TYPE}"
      echo ""
      echo "  如果启动失败，可以尝试以下故障排除方法:"
      echo "  1. 手动修复权限: docker exec -u root postgres-${SERVER_TYPE} chown -R postgres:postgres ${PG_DATA_DIR}"
      echo "  2. 检查容器日志: docker logs postgres-${SERVER_TYPE}"
      echo "  3. 以调试模式启动: docker exec -it postgres-${SERVER_TYPE} /usr/lib/postgresql/14/bin/postgres -D ${PG_DATA_DIR} -d 5"
    else
      echo "  ❌ 未找到有效的备份目录，请先传输备份"
      exit 1
    fi
  else
    echo "  ❌ 备份目录不存在: ${BACKUP_DIR}/backup"
    echo "  请先传输备份文件到此目录"
    exit 1
  fi
else
  # 初始化操作流程
  echo "10. 环境初始化完成"
  echo "  请将备份文件传输到 ${BACKUP_DIR}/backup 目录"
  echo "  传输完成后，可以使用以下命令执行恢复操作:"
  echo "  docker exec -e OPERATION=restore pgbackrest-${SERVER_TYPE} /scripts/init-pgbackrest-server-b-enhanced.sh"
  echo ""
  echo "  或手动执行恢复命令:"
  echo "  pgbackrest --stanza=${STANZA_NAME} --log-level-console=${LOG_LEVEL} restore"
fi

echo "===== pgBackRest 服务器B 环境初始化完成 ====="
echo ""
echo "配置摘要:"
echo "  备份目录: ${BACKUP_DIR}"
echo "  数据目录: ${PG_DATA_DIR}"
echo "  Stanza名称: ${STANZA_NAME}"
echo ""
echo "---------------------------------------------------"
echo "备份文件传输指南:"
echo "---------------------------------------------------"
echo "1. 从服务器A复制备份文件到服务器B:"
echo "   docker cp pgbackrest-A:${BACKUP_DIR}/archive/pgbackup_*.tar.gz ${BACKUP_DIR}/archive/"
echo "   # 或者使用数据卷复制:"
echo "   cp /path/to/backup/pgbackup_*.tar.gz ${BACKUP_DIR}/archive/"
echo ""
echo "2. 在服务器B上解压备份文件:"
echo "   tar -xzf ${BACKUP_DIR}/archive/pgbackup_*.tar.gz -C ${BACKUP_DIR}"
echo ""
echo "3. 执行恢复操作:"
echo "   docker exec -e OPERATION=restore pgbackrest-${SERVER_TYPE} /scripts/init-pgbackrest-server-b-enhanced.sh"
echo "---------------------------------------------------"
echo "注意: 请确保PostgreSQL服务已停止再执行恢复操作"
echo "---------------------------------------------------"
